"""
    create by IAmFiveHigh on 2024/8/24
"""
from enum import Enum


class RedBlack(Enum):
    RED = True
    BLACK = False


class RBTree:
    class Node:
        def __init__(self, key, value=None):
            self.key = key
            self.value = value
            self.left = None
            self.right = None
            self.color = RedBlack.RED
            self.height = 1

    def __init__(self):
        self._root = None
        self._size = 0

    def is_empty(self) -> bool:
        return self._size == 0

    def get_size(self) -> int:
        return self._size

    def add(self, key, value):
        if self.is_empty():
            self._root = RBTree.Node(key, value)
            self._size += 1
            self._root.color = RedBlack.BLACK
            return

        self._root = self._add(self._root, key, value)
        # 最终根节点一定为black
        self._root.color = RedBlack.BLACK

    def _add(self, node: Node, key, value):
        if node is None:
            # 如果传进来的node为空，用key value生成新的红色node
            self._size += 1
            return RBTree.Node(key, value)
        if key < node.key:
            node.left = self._add(node.left, key, value)
        elif key > node.key:
            node.right = self._add(node.right, key, value)
        else:
            node.value = value
        #     红黑树              23树                                红黑树             23树
        #
        #                                           左旋
        #  node          ===    node new(red)    ------->        new(red)       ===    node(red) new(red)
        #    \                                                  /
        #      new(red)                                        node(red)
        if self.is_red(node.right) and not self.is_red(node.left):
            node = self._left_rotate(node)
        #      红黑树                     23树                                      红黑树
        #       node                                                右旋            left
        #      /                                                                 /       \
        #    left(red)          ===     new(red) left(red) node    ------->    new(red)   node(red)
        #   /
        # new(red)
        if self.is_red(node.left) and node.left is not None and self.is_red(node.left.left):
            node = self._right_rotate(node)

        #       node                            node(red)
        #      /     \            ---->       /        \
        #  left(red)  right(red)         left(black)  right(black)
        if self.is_red(node.left) and self.is_red(node.right):
            self._flip_color(node)

        # 更新node的高度
        node.height = max(self._get_node_height(node.left), self._get_node_height(node.right)) + 1

        return node

    @staticmethod
    def is_red(node: Node):
        if node is None:
            return False
        return True if node.color is RedBlack.RED else False

    #     node                      x
    #     / \                    /  \
    #  T1     x      左旋转    node     T3
    #        /  \    ----->   / \
    #       T2   T3          T1  T2
    @staticmethod
    def _left_rotate(node):
        x = node.right
        node.right = x.left
        x.left = node

        # x替代了node的位置所以继承node的color
        x.color = node.color
        # node在旋转后在左边替代了x的新结点的身份，所以是红色
        node.color = RedBlack.RED

        return x

    #     node                      x
    #     / \                    /  \
    #    x     T1    右旋转    T3     node
    #  /  \          ----->          /  \
    # T3   T2                       T2   T1
    @staticmethod
    def _right_rotate(node):
        x = node.left
        node.left = x.right
        x.right = node

        # x替代了node的位置所以继承node的color
        x.color = node.color
        # 右旋转后左右两边和中间三个值在一个结点，所以左右两边都是红色
        node.color = RedBlack.RED

        return x

    # 颜色反转
    @staticmethod
    def _flip_color(node: Node):
        # 中间节点要向上和上面的合并，所以是红色
        node.color = RedBlack.RED
        # 子结点为黑
        node.left.color = RedBlack.BLACK
        node.right.color = RedBlack.BLACK

    # 判断是否为二分搜索树
    def is_bst(self):
        if self.is_empty():
            return False
        order_result = []
        self._in_order(self._root, lambda k, v: order_result.append(k))
        # 如果是二分搜索树，中序遍历的数组应该从小到大排序
        for index in range(1, len(order_result)):
            if order_result[index] < order_result[index - 1]:
                return False
        return True

    # 中序遍历
    def in_order(self, func):
        if self.is_empty():
            return
        self._in_order(self._root, func)

    def _in_order(self, node, func):
        if node is None:
            return

        self._in_order(node.left, func)
        func(node.key, node.value)
        self._in_order(node.right, func)

    @staticmethod
    def _get_node_height(node):
        if node is None:
            return 0
        return node.height

    def _get_balance(self, node):
        if node is None:
            return 0
        return self._get_node_height(node.left) - self._get_node_height(node.right)

    # 判断是否是平衡二叉树
    def is_balanced(self):
        return self._is_balanced(self._root)

    def _is_balanced(self, node):
        if node is None:
            return True
        balance_factor = self._get_balance(node)
        # 任意node的平衡因子不能大于1 否则就不算平衡二叉树
        if abs(balance_factor) > 1:
            return False

        return self._is_balanced(node.left) and self._is_balanced(node.right)


if __name__ == '__main__':

    def test_tree(t, result_list, error_message):
        in_order_list = []
        t.in_order(lambda k, v: in_order_list.append(k))
        if in_order_list != result_list:
            print(error_message)


    temp = [5, 6, 1, 3, 4, 2]
    tree = RBTree()
    for i in temp:
        tree.add(i, None)
    test_tree(tree, [1, 2, 3, 4, 5, 6], "二叉树添加元素方法有bug")

    if not tree.is_balanced():
        print("不是平衡二叉树")
